#include "disk.h"
#include "cpu_instr.h"
#include "log.h"
#include "klib.h"
#include "boot_info.h"
#include "dev.h"
#include "irq.h"
#include "mutex.h"
#include "sem.h"

static disk_t  disk_buf[DISK_CNT];
static sem_t   sem;
static mutex_t mutex;
static int     disk_on_flag = 0;

static void disk_send_cmd(disk_t *disk, uint32_t start_sector, uint32_t sector_count, int cmd)
{
    uint8_t disk_choose = DISK_DRIVE_BASE | disk->drive;
    outb(DISK_DRIVE(disk), disk_choose);

    outb(DISK_SECTOR_COUNT(disk), (uint8_t)(sector_count >> 8));
    outb(DISK_LBA_LO(disk), (uint8_t)(start_sector >> 24));
    outb(DISK_LBA_MID(disk), 0);
    outb(DISK_LBA_HI(disk), 0);

    outb(DISK_SECTOR_COUNT(disk), (uint8_t)(sector_count));
    outb(DISK_LBA_LO(disk), (uint8_t)start_sector);
    outb(DISK_LBA_MID(disk), (uint8_t)(start_sector >> 8));
    outb(DISK_LBA_HI(disk), (uint8_t)(start_sector >> 16));

    outb(DISK_CMD(disk), (uint8_t)cmd);
}
static void disk_read_data(disk_t *disk, void *buf, int size)
{
    uint16_t *c = (uint16_t *)buf;
    for (int i = 0; i < size / 2; i++)
    {
        *c++ = inw(DISK_DATA(disk));
    }
}
static void disk_write_data(disk_t *disk, void *buf, int size)
{
    uint16_t *c = (uint16_t *)buf;
    for (int i = 0; i < size / 2; i++)
    {
        outw(DISK_DATA(disk), *c++);
    }
}
static int disk_wait_data(disk_t *disk)
{
    uint8_t status;
    do
    {
        status = inb(DISK_STATUS(disk));
        if ((status & (DISK_STATUS_BUSY | DISK_STATUS_DRQ | DISK_STATUS_ERR)) != DISK_STATUS_BUSY)
        {
            break;
        }
    } while (1);
    return (status & DISK_STATUS_ERR) ? -1 : 0;
#if 0
    while((inb(DISK_STATUS(disk)) & DISK_STATUS_BUSY));
    if((inb(DISK_STATUS(disk)) & DISK_STATUS_ERR)){
        log_printf("ERR disk wait data, disk failed (ERR=1)");
        return -1;
    }
    while(!((inb(DISK_STATUS(disk)) & DISK_STATUS_DRQ)));
    return 0;
#endif
}
static int detect_part_info(disk_t *disk)
{
    mbr_t mbr;
    disk_send_cmd(disk, 0, 1, DISK_CMD_READ);
    int err = disk_wait_data(disk);
    if (err < 0)
    {
        log_printf("read mbr failed");
        return err;
    }
    disk_read_data(disk, &mbr, sizeof(mbr));
    // traversing all part info
    part_item_t *item      = mbr.part_item;
    partinfo_t  *part_info = disk->partinfo + 1;
    for (int i = 0; i < MBR_PRIMARY_PART_NR; i++, item++, part_info++)
    {
        part_info->type = item->system_id;
        if (part_info->type == FS_INVALID)
        {
            part_info->total_sector = 0;
            part_info->start_sector = 0;
            part_info->disk         = (disk_t *)0;
        }
        else
        {
            kernel_sprintf(part_info->name, "%s%d", disk->name, i + 1);
            part_info->start_sector = item->relative_sectors;
            part_info->total_sector = item->total_sectors;
            part_info->disk         = disk;
        }
    }
}
static int identify_disk(disk_t *disk)
{
    disk_send_cmd(disk, 0, 0, DISK_CMD_IDENTITY);
    int err = inb(DISK_STATUS(disk));
    if (err == 0)
    {
        log_printf("%s doesn't exist", disk->name);
        return -1;
    }
    err = disk_wait_data(disk);
    if (err < 0)
    {
        log_printf("disk[%s]: read failed", disk->name);
        return err;
    }
    uint16_t buf[256];
    disk_read_data(disk, buf, sizeof(buf));
    disk->sector_count = *(uint32_t *)(buf + 100);
    disk->sector_size  = SECTOR_SIZE;

    // diskpart info detect
    partinfo_t *part = disk->partinfo + 0;
    part->disk       = disk;
    kernel_sprintf(part->name, "%s%d", disk->name, 0);
    part->start_sector = 0;
    part->total_sector = disk->sector_count;
    part->type         = FS_INVALID;
    detect_part_info(disk);
    return 0;
}
static void print_disk_info(disk_t *disk)
{
    log_printf("    disk name:%s", disk->name);
    log_printf("    port base: %x", disk->port_base);
    log_printf("    total size:%d m", disk->sector_count * disk->sector_size / 1024 / (1024));
    log_printf("    drive:%s,", disk->drive == DRIVE_MASTER ? "Master" : "Slave");
    log_printf("    Part info:");
    for (int i = 0; i < DISK_PRIMARY_PART_CNT; i++)
    {
        partinfo_t *part_info = disk->partinfo + i;
        if (part_info->type != FS_INVALID)
        {
            log_printf("        part_info[%d]:%s: type:%x,start_sector:%d,count %d",
                       i,
                       part_info->name,
                       part_info->type,
                       part_info->start_sector,
                       part_info->total_sector);
        }
    }
}

void read_slave_test(void)
{
    disk_t disk;
    kernel_memset(&disk, 0, sizeof(disk_t));
    disk_t *disk_slave    = &disk;
    disk_slave->port_base = IOBASE_SECONDARY;
    outb(0x176, 0xF0);
    outb(0x177, 0xEC);
    while ((inb(0x177) & 0x80))
        ;
    if ((inb(0x177) & 0x01))
    {
        log_printf("ERR disk wait data, disk failed (ERR=1)");
    }
    while (!((inb(0x177) & 0x08)))
        ;
    uint16_t buf16[256];
    disk_read_data(disk_slave, buf16, sizeof(buf16));
}
void disk_init(void)
{
    // read_slave();
    log_printf("check disk......");
    kernel_memset(disk_buf, 0, sizeof(disk_buf));
    mutex_init(&mutex);
    sem_init(&sem, 0);
    for (int i = 0; i < DISK_PER_CHANNEL; i++)
    {
        disk_t *disk = disk_buf + i;
        kernel_sprintf(disk->name, "sd%c", i + 'a');
        disk->drive     = (i == 0) ? DRIVE_MASTER : DRIVE_SLAVE;
        disk->port_base = IOBASE_PRIMARY;
        int err         = identify_disk(disk);
        if (err == 0)
        {
            print_disk_info(disk);
        }
        disk->mutex = &mutex;
        disk->sem   = &sem;
    }
    for (int i = 0; i < DISK_PER_CHANNEL; i++)
    {
        disk_t *disk = disk_buf + 2 + i;
        kernel_sprintf(disk->name, "sd%c", i + 'c');
        disk->drive     = (i == 0) ? DRIVE_MASTER : DRIVE_SLAVE;
        disk->port_base = IOBASE_SECONDARY;
        int err         = identify_disk(disk);
        if (err == 0)
        {
            print_disk_info(disk);
        }
        disk->mutex = &mutex;
        disk->sem   = &sem;
    }
}
int disk_open(device_t *dev)
{
    // 0xa0,disk:a,b,c part:0,2,3
    int disk_idx = (dev->minor >> 4) - 0xa;    // a,b,c->0,1,2
    int part_idx = dev->minor & 0xF;
    if ((disk_idx >= DISK_CNT) || (part_idx >= DISK_PRIMARY_PART_CNT))
    {
        log_printf("disk device minor error:%d", dev->minor);
        return -1;
    }
    disk_t *disk = disk_buf + disk_idx;
    if (disk->sector_count == 0)
    {
        log_printf("disk not exist, device:sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    partinfo_t *part_info = disk->partinfo + part_idx;
    if (part_info->total_sector == 0)
    {
        log_printf("part not exit, device:sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    dev->data = (void *)part_info;
    irq_install(IRQ14_DISK_PRI, (irq_handler_t)exception_handler_disk_primary);
    irq_enable(IRQ14_DISK_PRI);
    return 0;
}
int disk_read(device_t *dev, int addr, char *buf, int size)
{                                              // the addr and size is sector, not bytes
    int disk_idx = (dev->minor >> 4) - 0xa;    // a,b,c->0,1,2
    int part_idx = dev->minor & 0xF;
    // size=up2(size,SECTOR_SIZE);
    partinfo_t *part_info = (partinfo_t *)dev->data;
    // send read req
    if (!part_info)
    {
        log_printf("disk_read:get part info failed. device: sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    disk_t *disk = part_info->disk;
    if (disk == (disk_t *)0)
    {
        log_printf("disk_read: the part info has no disk:device :sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    mutex_lock(disk->mutex);
    disk_on_flag = 1;
    disk_send_cmd(disk, part_info->start_sector + addr, size, DISK_CMD_READ);
    int cnt;
    for (cnt = 0; cnt < size; cnt++, buf += disk->sector_size)
    {
        if (task_current_task())
        {
            sem_wait(disk->sem);
        }
        int err = disk_wait_data(disk);
        if (err < 0)
        {
            log_printf("disk read:disk part_info(%s) read err: start sector %d,count:%d", part_info->name, addr, size);
            break;
        }
        // dbr_t dbr;
        // disk_read_data(disk,&dbr,disk->sector_size);
        disk_read_data(disk, buf, disk->sector_size);    // read 1 sector data
    }
    mutex_unlock(disk->mutex);
    disk_on_flag = 0;
    // wait
    // red
    return cnt;
}
int disk_write(device_t *dev, int addr, char *buf, int size)
{
    int disk_idx = (dev->minor >> 4) - 0xa;    // a,b,c->0,1,2
    int part_idx = dev->minor & 0xF;
    // size=up2(size,SECTOR_SIZE);
    partinfo_t *part_info = (partinfo_t *)dev->data;
    // send read req
    if (!part_info)
    {
        log_printf("disk_write:get part info failed. device: sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    disk_t *disk = part_info->disk;
    if (disk == (disk_t *)0)
    {
        log_printf("disk_write: the part info has no disk:device :sd%c%d", disk_idx + 'a', part_idx);
        return -1;
    }
    mutex_lock(disk->mutex);
    disk_on_flag = 1;
    disk_send_cmd(disk, part_info->start_sector + addr, size, DISK_CMD_WRITE);
    int cnt;
    for (cnt = 0; cnt < size; cnt++, buf++)
    {
        disk_write_data(disk, buf, disk->sector_size);
        int err = disk_wait_data(disk);
        if (err < 0)
        {
            log_printf("disk_write:disk part_info(%s) read err: start sector %d,count:%d", part_info->name, addr, size);
            break;
        }
        if (task_current_task())
        {
            sem_wait(disk->sem);
        }
    }
    mutex_unlock(disk->mutex);
    disk_on_flag = 0;
    // wait
    // red
    return cnt;
}
int disk_control(device_t *dev, int cmd, int arg0, int arg1)
{
    return -1;
}
void disk_close(device_t *dev)
{
    return;
}

void do_handler_disk_primary(exception_frame_t *frame)
{

    pic_send_eoi(IRQ14_DISK_PRI);
    if (disk_on_flag && task_current_task())
    {
        sem_notify(&sem);
    }
}

dev_desc_t dev_disk_desc = {
    .name    = "disk",
    .major   = DEV_DISK,
    .open    = disk_open,
    .read    = disk_read,
    .write   = disk_write,
    .control = disk_control,
    .close   = disk_close,
};